 # -*- coding: utf-8 -*-

from __future__ import print_function
import keras
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras import backend as K

import win_unicode_console
win_unicode_console.enable()

import tensorflow.compat.v1 as tf
tf.enable_eager_execution()

import time

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

def buildModel(input_shape,num_classes):
    model = Sequential()
    # 加入卷积层
    model.add(Conv2D(32,
                      activation='relu',
                      input_shape=input_shape,                 
                      kernel_size = (3,3)))
    # 池化层
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
    # 加入Dropout层防止过拟合，提升模型泛化能力
    model.add(Dropout(0.25)) 
    #添加一个卷积层，包含64个卷积和，每个卷积和仍为3*3
    model.add(Conv2D(64, kernel_size=(3,3), activation='relu'))
    #来一个池化层
    model.add(MaxPooling2D(pool_size=(2,2), strides=(2,2)))
    # 加入Dropout层防止过拟合，提升模型泛化能力
    model.add(Dropout(0.25)) 
    # 压平层，常用在从卷积层到全连接层的过渡
    model.add(Flatten())
    # 第一层紧密连接128神经元
    model.add(tf.keras.layers.Dense(1000, activation=tf.nn.relu))
    # 加入Dropout层防止过拟合，提升模型泛化能力
    model.add(Dropout(0.25)) 
    # 第二层分10 个类别
    model.add(tf.keras.layers.Dense(num_classes, activation=tf.nn.softmax))
    # -----------------------compile----------------------------
    model.compile(loss=keras.metrics.categorical_crossentropy,
                  optimizer=keras.optimizers.RMSprop(lr=0.001, rho=0.9),
                  metrics=['accuracy'])
    model.summary()
    # 可参考 https://blog.csdn.net/qq_35200479/article/details/83752487
    return model
    pass

# 保存db模型
def saveModel():
    full_model = tf.function(lambda x: model(x))
    full_model = full_model.get_concrete_function(
        tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
    # Get frozen ConcreteFunction
    frozen_func = convert_variables_to_constants_v2(full_model)
    # print(frozen_func.graph.as_graph_def())
    tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                      logdir="",
                      name="model/keras_model.pb",
                      as_text=False)
    pass


batch_size = 128
num_classes = 10
epochs = 100
# input image dimensions
img_rows, img_cols = 28, 28
# the data, shuffled and split between train and test sets

from tensorflow.examples.tutorials.mnist import input_data  
MNIST_data_folder = "data"
fashion_mnist = input_data.read_data_sets(MNIST_data_folder,one_hot=True)

x_train, y_train = fashion_mnist.train.images,fashion_mnist.train.labels  
x_test, y_test = fashion_mnist.test.images, fashion_mnist.test.labels  
x_train = x_train.reshape(-1, 28, 28,1).astype('float32')  
x_test = x_test.reshape(-1,28, 28,1).astype('float32')

if K.image_data_format() == 'channels_first':
    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
    input_shape = (1, img_rows, img_cols)
else:
    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
    input_shape = (img_rows, img_cols, 1)

# 数据的归一化处理
x_train /= 255
x_test /= 255
print('x_train shape:', x_train.shape)
print('x_test shape:', x_test.shape)
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')

print('y_train shape:', y_train.shape)
print('y_test shape:', y_test.shape)
print(y_train.shape[0], 'train samples')
print(y_test.shape[0], 'test samples')


model = buildModel(input_shape,num_classes)

print('x_train shape:', x_train.shape)
print('x_test shape:', x_test.shape)
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')

start = time.time()
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs,
          verbose=1, validation_data=(x_test, y_test))
score = model.evaluate(x_test, y_test, verbose=0)

print('Test loss:', score[0])
print('Test accuracy:', score[1])
print('Cost %.1f s' % (time.time()- start))

saveModel()